// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.

import { ServiceBusMessage, toAmqpMessage, isServiceBusMessage } from "./serviceBusMessage";
import { throwTypeErrorIfParameterMissing } from "./util/errors";
import { ClientEntityContext } from "./clientEntityContext";
import {
  MessageAnnotations,
  messageProperties as RheaMessagePropertiesList,
  message as RheaMessageUtil
} from "rhea-promise";
import { AmqpMessage } from "@azure/core-amqp";

/**
 * The amount of bytes to reserve as overhead for a small message.
 */
const smallMessageOverhead = 5;
/**
 * The amount of bytes to reserve as overhead for a large message.
 */
const largeMessageOverhead = 8;
/**
 * The maximum number of bytes that a message may be to be considered small.
 */
const smallMessageMaxBytes = 255;

/**
 * A batch of messages that you can create using the {@link createBatch} method.
 *
 * @export
 */
export interface ServiceBusMessageBatch {
  /**
   * Size of the batch in bytes after the events added to it have been encoded into a single AMQP
   * message.
   * @readonly
   */
  readonly sizeInBytes: number;

  /**
   * Number of messages added to the batch.
   * @readonly
   */
  readonly count: number;

  /**
   * The maximum size of the batch, in bytes. The `tryAdd` function on the batch will return `false`
   * if the message being added causes the size of the batch to exceed this limit. Use the `createBatch()` method on
   * the `Sender` to set the maxSizeInBytes.
   * @readonly.
   */
  readonly maxSizeInBytes: number;

  /**
   * Adds a message to the batch if permitted by the batch's size limit.
   * **NOTE**: Always remember to check the return value of this method, before calling it again
   * for the next event.
   *
   * @param message  An individual service bus message.
   * @returns A boolean value indicating if the message has been added to the batch or not.
   */
  tryAdd(message: ServiceBusMessage): boolean;

  /**
   * The AMQP message containing encoded events that were added to the batch.
   * Used internally by the `sendBatch()` method on the `Sender`.
   * This is not meant for the user to use directly.
   *
   * @readonly
   * @internal
   * @ignore
   */
  _generateMessage(): Buffer;
}

/**
 * An internal class representing a batch of messages which can be used to send messages to Service Bus.
 *
 * @class
 * @internal
 * @ignore
 */
export class ServiceBusMessageBatchImpl implements ServiceBusMessageBatch {
  /**
   * @property Describes the amqp connection context for the Client.
   */
  private _context: ClientEntityContext;
  /**
   * @property The maximum size allowed for the batch.
   */
  private _maxSizeInBytes: number;
  /**
   * @property Current size of the batch in bytes.
   */
  private _sizeInBytes: number;
  /**
   * @property Encoded amqp messages.
   */
  private _encodedMessages: Buffer[] = [];
  /**
   * ServiceBusMessageBatch should not be constructed using `new ServiceBusMessageBatch()`
   * Use the `createBatch()` method on your `Sender` instead.
   * @constructor
   * @internal
   * @ignore
   */
  constructor(context: ClientEntityContext, maxSizeInBytes: number) {
    this._context = context;
    this._maxSizeInBytes = maxSizeInBytes;
    this._sizeInBytes = 0;
    this._batchMessageProperties = {};
  }

  /**
   * @property The maximum size of the batch, in bytes.
   * @readonly
   */
  get maxSizeInBytes(): number {
    return this._maxSizeInBytes;
  }

  /**
   * @property Size of the `ServiceBusMessageBatch` instance after the messages added to it have been
   * encoded into a single AMQP message.
   * @readonly
   */
  get sizeInBytes(): number {
    return this._sizeInBytes;
  }

  /**
   * @property Number of messages in the `ServiceBusMessageBatch` instance.
   * @readonly
   */
  get count(): number {
    return this._encodedMessages.length;
  }

  /**
   * Generates an AMQP message that contains the provided encoded messages and annotations.
   *
   * @private
   * @param {Buffer[]} encodedMessages The already encoded messages to include in the AMQP batch.
   * @param {MessageAnnotations} [annotations] The message annotations to set on the batch.
   * @param {{ [key: string]: any }} [applicationProperties] The application properties to set on the batch.
   * @param {{ [key: string]: string }} [messageProperties] The message properties to set on the batch.
   * @returns {Buffer}
   * @memberof ServiceBusMessageBatchImpl
   */
  private _generateBatch(
    encodedMessages: Buffer[],
    annotations?: MessageAnnotations,
    applicationProperties?: { [key: string]: any },
    messageProperties?: { [key: string]: string }
  ): Buffer {
    const batchEnvelope: AmqpMessage = {
      body: RheaMessageUtil.data_sections(encodedMessages),
      message_annotations: annotations,
      application_properties: applicationProperties
    };
    if (messageProperties) {
      for (const prop of RheaMessagePropertiesList) {
        if ((messageProperties as any)[prop]) {
          (batchEnvelope as any)[prop] = (messageProperties as any)[prop];
        }
      }
    }
    return RheaMessageUtil.encode(batchEnvelope);
  }

  /**
   * @property Represents the single AMQP message which is the result of encoding all the events
   * added into the `ServiceBusMessageBatch` instance.
   *
   * This is not meant for the user to use directly.
   *
   * When the `ServiceBusMessageBatch` instance is passed to the `sendBatch()` method on the `Sender`,
   * this single batched AMQP message is what gets sent over the wire to the service.
   * @readonly
   */
  _generateMessage(): Buffer {
    return this._generateBatch(
      this._encodedMessages,
      this._batchAnnotations,
      this._batchApplicationProperties,
      this._batchMessageProperties
    );
  }

  /**
   * The message annotations to apply on the batch envelope.
   * This will reflect the message annotations on the first message
   * that was added to the batch.
   */
  private _batchAnnotations?: MessageAnnotations;
  /**
   * The message properties to apply on the batch envelope.
   * This will reflect the message properties on the first message
   * that was added to the batch.
   */
  private _batchMessageProperties?: { [key: string]: string };
  /**
   * The application properties to apply on the batch envelope.
   * This will reflect the application properties on the first message
   * that was added to the batch.
   */
  private _batchApplicationProperties?: { [key: string]: any };

  /**
   * Tries to add a message to the batch if permitted by the batch's size limit.
   * **NOTE**: Always remember to check the return value of this method, before calling it again
   * for the next message.
   *
   * @param message  An individual service bus message.
   * @returns A boolean value indicating if the message has been added to the batch or not.
   */
  public tryAdd(message: ServiceBusMessage): boolean {
    throwTypeErrorIfParameterMissing(this._context.namespace.connectionId, "message", message);
    if (!isServiceBusMessage(message)) {
      throw new TypeError("Provided value for 'message' must be of type ServiceBusMessage.");
    }

    // Convert ServiceBusMessage to AmqpMessage.
    const amqpMessage = toAmqpMessage(message);
    amqpMessage.body = this._context.namespace.dataTransformer.encode(message.body);
    const encodedMessage = RheaMessageUtil.encode(amqpMessage);

    let currentSize = this._sizeInBytes;

    // The first time an event is added, we need to calculate
    // the overhead of creating an AMQP batch, including the
    // message_annotations, application_properties and message_properties
    // that are taken from the 1st message.
    if (this.count === 0) {
      if (amqpMessage.message_annotations) {
        this._batchAnnotations = amqpMessage.message_annotations;
      }
      if (amqpMessage.application_properties) {
        this._batchApplicationProperties = amqpMessage.application_properties;
      }
      for (const prop of RheaMessagePropertiesList) {
        if ((amqpMessage as any)[prop]) {
          (this._batchMessageProperties as any)[prop] = (amqpMessage as any)[prop];
        }
      }
      // Figure out the overhead of creating a batch by generating an empty batch
      // with the expected batch annotations.
      currentSize += this._generateBatch(
        [],
        this._batchAnnotations,
        this._batchApplicationProperties,
        this._batchMessageProperties
      ).length;
    }

    const messageSize = encodedMessage.length;
    const messageOverhead =
      messageSize <= smallMessageMaxBytes ? smallMessageOverhead : largeMessageOverhead;
    currentSize += messageSize + messageOverhead;
    // Check if the size of the batch exceeds the maximum allowed size
    // once we add the new event to it.
    if (currentSize > this._maxSizeInBytes) {
      return false;
    }

    // The message will fit in the batch, so it is now safe to store it.
    this._encodedMessages.push(encodedMessage);

    this._sizeInBytes = currentSize;
    return true;
  }
}
